# 训练一个模型： 输入一个5维向量arr  arr[0] - arr[3] >= arr[4] 为 1 否则为0
# 准备训练数据
import torch
import numpy as np
import torch.nn as nn

def train_data():
  x = np.random.randn(5)
  return x,np.argmax(x)

def gen_train(size):
  X = []
  Y = []
  for i in range(size):
    x,y=train_data()
    X.append(x)
    Y.append([y])
  return torch.FloatTensor(X),torch.LongTensor(Y)

class TorchModel(nn.Module):
  def __init__(self,inputSize=5,outPutSize=5):
    super(TorchModel,self).__init__()
    self.learn = nn.Linear(inputSize,outPutSize)
    # self.activity = nn.Sigmoid()
    self.loss = nn.functional.cross_entropy
  
  def forward(self,x_true,y_true=None):
    yp = self.learn(x_true)
    # yp = self.activity(yp)
    if y_true == None :
      return yp
    else:
      y_true = y_true.squeeze(1)
      return self.loss(yp,y_true)
def evaluate(model):
    model.eval() # 开启训练模式
    test_sample_num = 100 # 用于测试的数据量
    test_X,test_Y = gen_train(test_sample_num)
    successcount = 0
    with torch.no_grad():
        for x,y_true in zip(test_X, test_Y):
            y_prev = model(x)
            if np.argmax(y_prev) == y_true:
              successcount+= 1
    bili = successcount / test_sample_num
    print(f"本轮训练结束,模型的正确率{bili:.3f}")
    return bili
def main():
  batch_size=10 # 100个更新一次权重
  lr=0.0001 # 学习率
  eposhs=20 # 跑20轮
  train_dateNum=10000 # 训练集的数据量
  model = TorchModel(5,5)
  adam=torch.optim.Adam(model.parameters(),lr=lr)
  X,Y = gen_train(train_dateNum)
  log = []
  for epoch in range(eposhs):
    model.train()
    watchLoss = []
    for batch in range(int(train_dateNum/batch_size)):
      x_test = X[batch*batch_size:(batch+1)*batch_size,:]
      y_test = Y[batch*batch_size:(batch+1)*batch_size,:]
      loss = model(x_test,y_test)
      loss.backward() # 计算梯度
      adam.step() # 更新梯度
      adam.zero_grad() # 梯度归零
      watchLoss.append(loss.item())
    # 每一轮结束测试模型
    print(f"第{epoch}轮 , 平均损失值为{np.mean(watchLoss):.5f}") ## 试试这个语法
    # 每轮训练完都测试一下这一轮的成功率
    acc = evaluate(model)
    log.append([acc, float(np.mean(watchLoss))])
  torch.save(model.state_dict(),'model.ini')
  print(log)
if __name__ == "__main__":
  main()